{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Prune the model\n",
    "In this step, we will explore two methods to prune the model - depth and width pruning. Refer to the [README.md](./README.md) to decide which pruning techniques you would like to explore. For usage details, please refer to the [pruning docs](https://docs.nvidia.com/nemo-framework/user-guide/latest/model-optimization/pruning/pruning.html) for more details.\n",
    "\n",
    "Let's define the common parameters for depth or width pruning first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NEMO_ROOT = \"/opt/NeMo\"\n",
    "ROOT_DIR = \"/workspace\"\n",
    "MODEL_PATH = f\"{ROOT_DIR}/Qwen3-8B-nemo\"\n",
    "\n",
    "##### Set data paths\n",
    "# NOTE: If you have multiple partitioned datasets, you can pass in a space-separated list of paths below.\n",
    "DATA_PATH = f\"{ROOT_DIR}/wikitext-data\"\n",
    "DATA_PATHS = f\"{DATA_PATH}/wikitext-train_text_document\"\n",
    "INDEX_MAPPING_DIR = f\"{DATA_PATH}/index_mappings\"\n",
    "\n",
    "##### Set sequence length for pruning and distillation\n",
    "# NOTE: Use 4096 or 8192 depending on whether your dataset texts are short or long\n",
    "SEQ_LENGTH = 4096\n",
    "\n",
    "##### Change these to accommodate resources:\n",
    "# NOTE: Pruning only supports Tensor Parallelism (TP) 1. Number of layers in your model should be divisible by\n",
    "#   Pipeline Parallelism (PP) size, otherwise you can configure uneven PP using `--num_layers_in_first_pipeline_stage`\n",
    "#   and `--num_layers_in_last_pipeline_stage` arguments below with the gpt_prune.py script.\n",
    "DEVICES = 2\n",
    "TENSOR_PARALLEL_SIZE = 1\n",
    "PIPELINE_PARALLEL_SIZE = DEVICES\n",
    "MICRO_BATCH_SIZE = 4\n",
    "\n",
    "# Reduce this number to speed up the pruning process but may result in a slightly worse pruned model\n",
    "# Not used if directly dropping layers using `--drop_layers` argument.\n",
    "NUM_TRAIN_SAMPLES = 1024"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 2a: Using depth-pruning \n",
    "To depth-prune, we will prune the Qwen3-8B model from 36 to 24 layers resulting in a 6B model automatically by selecting the best 24 layers to keep based on activation statistics collected from the training samples.\n",
    "\n",
    "Alternatively, you can also directly drop layers 24-35 (1-indexed) using the `--drop_layers 24 25 26 27 28 29 30 31 32 33 34 35` argument (leaving 1-23 and 36) in the model which also works well generally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SAVE_PATH = f\"{ROOT_DIR}/Qwen3-8B-nemo-depth-pruned\"\n",
    "\n",
    "!torchrun --nproc_per_node \"{DEVICES}\" \"{NEMO_ROOT}/scripts/llm/gpt_prune.py\" \\\n",
    "    --devices \"{DEVICES}\" \\\n",
    "    --tp_size \"{TENSOR_PARALLEL_SIZE}\" \\\n",
    "    --pp_size \"{PIPELINE_PARALLEL_SIZE}\" \\\n",
    "    --restore_path \"{MODEL_PATH}\" \\\n",
    "    --legacy_ckpt \\\n",
    "    --save_path \"{SAVE_PATH}\" \\\n",
    "    --seq_length \"{SEQ_LENGTH}\" \\\n",
    "    --num_train_samples \"{NUM_TRAIN_SAMPLES}\" \\\n",
    "    --mbs \"{MICRO_BATCH_SIZE}\" \\\n",
    "    --data_paths \"{DATA_PATHS}\" \\\n",
    "    --index_mapping_dir \"{INDEX_MAPPING_DIR}\" \\\n",
    "    --target_num_layers 24"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running this script will save the depth-pruned model to your workspace at `<ROOT_DIR>/Qwen3-8B-nemo-depth-pruned`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 2b: Using width-pruning \n",
    "To width-prune, we will trim the `ffn_hidden_size` from 12288 to 9216 and `hidden_size` 4096 to 3584 also resulting in a 6B model. We can also trim the `num_attention_heads` and `num_query_groups` if needed. If the model is a Hybrid Mamba-Transformer model (e.g. [NVIDIA-Nemotron-Nano-12B-v2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2)), you can also trim the `mamba_num_heads` and `mamba_head_dim` dimensions.\n",
    "\n",
    "> **NOTE:** Pruning will take less then 10 minutes to run (depends on GPU)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SAVE_PATH = f\"{ROOT_DIR}/Qwen3-8B-nemo-width-pruned\"\n",
    "\n",
    "!torchrun --nproc_per_node \"{DEVICES}\" \"{NEMO_ROOT}/scripts/llm/gpt_prune.py\" \\\n",
    "    --devices \"{DEVICES}\" \\\n",
    "    --tp_size \"{TENSOR_PARALLEL_SIZE}\" \\\n",
    "    --pp_size \"{PIPELINE_PARALLEL_SIZE}\" \\\n",
    "    --restore_path \"{MODEL_PATH}\" \\\n",
    "    --legacy_ckpt \\\n",
    "    --save_path \"{SAVE_PATH}\" \\\n",
    "    --seq_length \"{SEQ_LENGTH}\" \\\n",
    "    --num_train_samples \"{NUM_TRAIN_SAMPLES}\" \\\n",
    "    --mbs \"{MICRO_BATCH_SIZE}\" \\\n",
    "    --data_paths \"{DATA_PATHS}\" \\\n",
    "    --index_mapping_dir \"{INDEX_MAPPING_DIR}\" \\\n",
    "    --target_ffn_hidden_size 9216 \\\n",
    "    --target_hidden_size 3584"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running this script will save the width-pruned model to your workspace at `<ROOT_DIR>/Qwen3-8B-nemo-width-pruned`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the depth and width pruned models, we can distill them from the unpruned model in next step."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
